from torch.utils.data import DataLoader, Dataset
import random
import linecache
import torch
import numpy as np
from PIL import Image
# 自定义数据集类，用于生成对比学习所需的图像对
class MyDataset(Dataset):
    def __init__(self,path_file,transform=None,should_invert=False):	# path_file是所有人脸图片的地址，每行地址是一个图片
        self.transform=transform
        self.should_invert=should_invert
        self.path_file = path_file

    def __getitem__(self, index):
        line=linecache.getline(self.path_file,random.randint(1,self.__len__()))
        img0_list=line.split("\\")
        #若为0，取得不同人的图片
        shouled_get_same_class=random.randint(0,1)
        if shouled_get_same_class:
            while True:
                img1_list=linecache.getline(self.path_file,random.randint(1,self.__len__())).split('\\')
                if img0_list[-1]==img1_list[-1]:
                   break

        else:
            while True:
                img1_list=linecache.getline(self.path_file,random.randint(1,self.__len__())).split('\\')
                if img0_list[-1]!=img1_list[-1]:
                    break

        img0_path = "/".join(img0_list).replace("\n","")
        img1_path = "/".join(img1_list).replace("\n","")

        im0=Image.open(img0_path).convert('L')
        im1=Image.open(img1_path).convert('L')


        im0 = torch.tensor(np.array(im0))
        im1 = torch.tensor(np.array(im1))

        return im0,im1,torch.tensor(shouled_get_same_class,dtype=torch.float32)

    def __len__(self):
        fh=open(self.path_file,'r')
        num=len(fh.readlines())
        fh.close()
        return num

import torch
from torch.utils.data import DataLoader
import demo2

device = "cuda"
device = "cpu"

net=demo2.SiameseNetwork().to(device)

criterion=demo2.ContrastiveLoss()
optimizer=torch.optim.Adam(net.parameters(),lr=0.001)

counter=[]
loss_history=[]
iteration_number=0

batch_size = 2
path_file = "/home/wangwei83/Desktop/pytorch-dl-cv/dataset/lfw/lfw-path_file.txt"
train_dataset = MyDataset(path_file=path_file)
train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,num_workers=0,pin_memory=True)

for epoch in range(0,20):
    for i,data in enumerate(train_loader,0):
        img0,img1,label=data
        img0,img1,label=img0.float().to(device),img1.float().to(device),label.to(device)
        optimizer.zero_grad()
        output1,output2=net(img0,img1)

        loss_contrastive=criterion(output1,output2,label)
        loss_contrastive.backward()
        optimizer.step()

        if i % 2 ==0:
            print('epoch:{},loss:{}\n'.format(epoch,loss_contrastive.item()))
            counter.append(iteration_number)
            loss_history.append(loss_contrastive.item())


